package com.lglbc.langchain4jdemo;

import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class MongoDBChatMemoryStore implements ChatMemoryStore {
    @Autowired
    private MongoTemplate mongoTemplate;
    @Override
    public List<dev.langchain4j.data.message.ChatMessage> getMessages(Object o) {
        Criteria memoryId = Criteria.where("memoryId").is(o);
        com.lglbc.langchain4jdemo.ChatMessage message = mongoTemplate.findOne(Query.query(memoryId), com.lglbc.langchain4jdemo.ChatMessage.class);
        if (message == null || message.getContent() == null) {
            return new ArrayList<>();
        }
        return ChatMessageDeserializer.messagesFromJson(message.getContent());
    }

    @Override
    public void updateMessages(Object o, List<dev.langchain4j.data.message.ChatMessage> list) {
        Criteria memoryId = Criteria.where("memoryId").is(o);
        Update update = new Update();
        update.set("content", ChatMessageSerializer.messagesToJson(list));
        mongoTemplate.upsert(Query.query(memoryId), update, com.lglbc.langchain4jdemo.ChatMessage.class);
    }

    @Override
    public void deleteMessages(Object o) {
        Criteria memoryId = Criteria.where("memoryId").is(o);
        mongoTemplate.remove(Query.query(memoryId), com.lglbc.langchain4jdemo.ChatMessage.class);
    }
}
